//////////////////////////////////////////////////////////////////
// MattiseOS Kernel
// Date: 2008
// Author: Matthew Iselin
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// at your option) any later version.
// 
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
// 
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.
//////////////////////////////////////////////////////////////////

// Includes
#include <mattise.h>
#include <stdint.h>
#include <mem/util.h>
#include <mem/heap.h>
#include <mem/page.h>
#include <kernio/io.h>
#include <process/common.h>

// DLMALLOC SUPPORT ROUTINES

uint32_t current_kheap_base = 0;
uint32_t next_kheap_base = 0;
#define HEAP_PAGE_FLAGS (PF_PRESENT | PF_READWRITE | PF_KERNELPAGE | PF_IMMUTABLE | PF_USER)

extern int paging_ready_to_go;
extern int pagefault_just_continue_allowed;

void* dlmalloc_extend_heap(int32_t incr)
{
	dprintf("%s(%d)\n", __FUNCTION__, incr);

	// check if we need to lose the space
	if(incr < 0)
	{
		dprintf("need to deallocate space\n");
		while(1);
	}

	// sbrk(0) finds the *current* location
	if(incr == 0)
		return (void*) current_kheap_base;

	// set up the task
	if( current_kheap_base == 0 )
	{
		// setup the heap
		current_kheap_base = next_kheap_base = KERNEL_HEAP_BASE;
		uint32_t pagecount = (incr / 0x1000);
		if(pagecount == 0)
			pagecount = 1;
		kAllocateAndMapPageDir( GetKernelDirectory(), pagecount, current_kheap_base, HEAP_PAGE_FLAGS );
		next_kheap_base += incr;

		// return it
		return (void*) KERNEL_HEAP_BASE;
	}

	// here we go...

	// need to allocate an extra (incr) bytes
	uint32_t newheap_base = next_kheap_base;
	uint32_t newheap_top = (next_kheap_base += incr);

	// find out how many pages that crosses (dlmalloc ALWAYS allocates at least a page)
	uint32_t pagecount = (newheap_top - newheap_base) / 0x1000;

	// do the deed
	kAllocateAndMapPageDir( GetKernelDirectory(), pagecount, newheap_base, HEAP_PAGE_FLAGS );

	// set the new current location
	current_kheap_base = newheap_base;

	// save the old heap location and get ready
	/*uint32_t oldheap = current_kheap_base;
	current_kheap_base = next_kheap_base;
	next_kheap_base += incr;

	// detect a page skip (if none we don't need to alloacte any more pages)
	uint32_t page1 = oldheap & ~0xFFF;
	uint32_t page2 = next_kheap_base & ~0xFFF;
	if( page1 != page2 )
	{
		// do it
		uint32_t pagecount = (page2-page1) / 0x1000;
		kAllocateAndMapPageDir( GetKernelDirectory(), pagecount, page1, HEAP_PAGE_FLAGS );
	}*/

	// return the old top of the heap
	return (void*) newheap_base;
}

void dlmalloc_abort(const char* source)
{
	dprintf("[MEM  ] dlmalloc_abort [%s]...\n", source);
	*((uint32_t*) 0) = 0xdeadbeef;
	kill(getpid());
}

#if 0

// Types

// a block of memory in the address space
struct memblock
{
	uint32_t used; // 4 bytes
	uint32_t size; // 4 bytes
	struct memblock* next; // 4 bytes
} __attribute__((packed)); // 12 bytes

// Variables

// the end of the executable image (the address of the variable, not the value)
extern char __end;

// the kernel heap
static struct memblock* kHeap = 0;

// the current virtual base of the kernel heap
uint32_t kheap_base;

// the size of the kernel heap initially
#define SZ_KERNEL_HEAP 1

// Functions

// allocates memory on the heap
void* malloc( size_t size )
{
	// locals that will be used in the allocation
	struct memblock* mblock, *newblock, *prev;
	char* cblock;
	size_t size2;

	// check for a nothing call
	if( size == 0 )
	{
		dprintf("== [WARN ] malloc with null size ==\n");
		return 0;
	}

	// calculate the actual size that the block will require
	size2 = size + sizeof( struct memblock );

	// check that there is already a heap
	if( kHeap == 0 )
	{
		// set the kernel heap location
		kheap_base = 0xb0000000;
		kHeap = mblock = (struct memblock*) kAllocateAndMap(SZ_KERNEL_HEAP, kheap_base, PF_PRESENT | PF_KERNELPAGE | PF_USER | PF_READWRITE | PF_IMMUTABLE);
		kheap_base += SZ_KERNEL_HEAP * 0x1000;

		// setup it's state
		mblock->used = 0;
		mblock->size = SZ_KERNEL_HEAP * 0x1000;
		mblock->next = 0;
	}

	// look for a block (effectively this provides us with a linked list
	// of blocks with each entry being a header to a block of memory)
	for( mblock = kHeap; mblock != 0; prev = mblock, mblock = mblock->next )
		if( (!mblock->used) && (size2 <= mblock->size) )
			break;

	// check that one was found
	if( mblock == 0 )
	{
		// check for a valid prev variable (needed to work properly)
		if(prev != 0)
		{
			// attempt to extend the heap by two pages (or more, depending on size)
			// by using two pages we know that we'll always be able to fit in a bit
			// more later without having to map more pages, which is a slow process
			uint32_t extend_sz = size > 0x1000 ? ((size & ~0xFFF) / 0x1000) + 1 : 1;
			uint32_t old_base = kheap_base;
			dprintf("[MEM  ] extending heap by %d pages [size=%d, old=%x]... ", extend_sz, size, old_base);
			kheap_base += extend_sz * 0x1000;
			struct memblock* b = (struct memblock*) kAllocateAndMap(extend_sz, old_base, PF_PRESENT | PF_KERNELPAGE | PF_USER | PF_READWRITE | PF_IMMUTABLE);
			dprintf("Done, now at %x!\n", b);
			if(b == 0)
				mblock = 0;
			else
			{
				// link it in (merge if needed)
				if(prev->used == 0)
					prev->size += (extend_sz * 0x1000); // the final byte is not mapped
				else
				{
					prev->next = b;
					b->next = 0;
					b->size = (extend_sz * 0x1000);
					b->used = 0;
				}

				// search again
				for( mblock = kHeap; mblock != 0; mblock = mblock->next )
					if( (!mblock->used) && (size2 <= mblock->size) )
						break;
			}
		}

		// if there still isn't a valid block, give up
		if(mblock == 0)
		{
			dprintf("[MEM  ] malloc returns null\n");
			return 0;
		}
	}

	// setup the block (cblock holds the entire block)
	cblock = (char*) mblock + size2;
	newblock = (struct memblock*) cblock;

	dprintf("mblock = %x, size of memblock = %x, newblock = %x\n", mblock, sizeof(struct memblock), newblock);

	// point this entry to a new block with our settings
	newblock->used = 0;
	newblock->size = mblock->size - size2;
	newblock->next = mblock->next;

	asm volatile("nop; nop; nop");

	// point the old block to the new block and set it up
	mblock->used = 1;
	mblock->size = size;
	mblock->next = newblock;

	// return a pointer to the memory
	cblock = (char*) mblock;
	return (void*) (cblock + sizeof( struct memblock ));
}

// frees memory from the heap
void free( void* ptr )
{
	// locals for use later
	struct memblock* mblock;
	char* cblock;

	// get the real address of the block to free
	cblock = (char*) ptr - sizeof( struct memblock );

	// find this address in the linked list
	for( mblock = kHeap; mblock != 0; mblock = mblock->next )
	{
		// check that it's equal
		if( mblock == (struct memblock*) cblock )
			break;
	}

	// check that the pointer is a valid entry
	if( mblock == 0 )
		return;

	// free this block
	mblock->used = 0;

	//return;

	// try to merge unused blocks together
	for(mblock = kHeap; mblock != 0; mblock = mblock->next)
	{
		// merge as many as we can in one go
		while(1)
		{
			// currently we have a block, check that it's used
			if(mblock->used == 0)
			{
				// no, check that the next is valid and that it's not used
				if(mblock->next)
				{
					if(mblock->next->used == 0)
					{
						// it's not used, so fiddle with the size
						mblock->size += mblock->next->size + sizeof(struct memblock);
						mblock->next = mblock->next->next;
					}
					else
						break;
				}
				else
					break;
			}
			else
				break;
		}
		// look for valid blocks to merge
		//while( ( !mblock->used ) && ( mblock->next != 0 ) && ( !mblock->next->used ) )
		//{
			// merge the two blocks (their size should add up and the pointers should change)
		//	mblock->size += mblock->next->size + sizeof(struct memblock); // ((uint32_t) mblock->next) - ((uint32_t) mblock) + mblock->next->size;
		//	mblock->next = mblock->next->next;
		//}
	}
}

// reallocates memory
void* realloc( void* ptr, size_t newsz )
{
	// check the given pointer
	if( ptr == 0 )
		return 0;

	// try to allocate a new block
	void* newblock = malloc( newsz );

	// if it's invalid give up
	if( newblock == 0 )
		return 0;

	// get the memory block information for the old pointer
	struct memblock* oldinfo = (struct memblock*) (((uint32_t) ptr) - sizeof( struct memblock ));

	// it's valid, copy the old data into this block
	memcpy( newblock, ptr, (newsz > oldinfo->size ? oldinfo->size : newsz) );

	// free the old pointer
	free( ptr );

	// and now, return the pointer
	return newblock;
}

#endif
